import paramiko
import time
import os
import sys
import random


class SshToClient():

    """连接到服务器并执行相关动作
    Args:
        host:       远程服务器地址,ip.
        port:       远程服务器登录端口
        user:       远程登录用户, 非root
        user_pwd:   登录用户的密码
        root_pwd:   root密码
    """

    cmd_default    = 'uptime'
    cmd_root_login = 'su - root'

    def __init__(self,host,port,user=None,user_pwd=None,root_pwd=None):
        self.host     = host
        self.port     = port
        self.user     = user
        self.user_pwd = user_pwd
        self.root_pwd = root_pwd
        # self.cmd_default='uptime'
        # self.cmd_root_login='su - root'

    def exec_command_root(self,cmds=[cmd_default]):

        """切换到root用户执行多个命令
        Args:
            cmds:   待执行的命令列表,默认为['uptime'].
        """

        print('即将连接到%s' % self.host)
        print('@@@@@@@@@@@@连接成功@@@@@@@@@@@@')
        s = paramiko.SSHClient()
        s.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            s.connect(self.host, self.port, self.user, self.user_pwd,timeout=30)

            if self.user == 'root':
                # 运行命令
                print('使用root登录')
                results = []
                for cmd in cmds:
                    stdin,stdout,stderr = s.exec_command(cmd)
                    results.append(stdout.read().decode('gbk'))
                s.close

                for res in results: print(res)
                return results
            else:
                print(f'使用{self.user}登录')
                #启动子shell
                ssh = s.invoke_shell()
                time.sleep(0.1)
                buff = ssh.recv(9999).decode('gbk')
                print(buff,end='')

                #执行su - root切换到root用户
                # print(f'切换root登录')
                # result = buff
                result = ''
                buff=''
                ssh.send(self.cmd_root_login+'\n')
                buff = ssh.recv(9999).decode('gbk')
                while not buff.rstrip().endswith('Password:'):
                    buff += ssh.recv(9999).decode('gbk')
                # print(buff) 
                
                #输入root用户密码
                result += buff
                buff=''
                ssh.send(self.root_pwd+'\n')
                buff = ssh.recv(9999).decode('gbk')
                while not buff.rstrip().endswith('#'):
                    buff += ssh.recv(9999).decode('gbk')
                    if buff.rstrip().endswith('$') or buff.find('denied')>=0:
                        print('当前用户无法切换到root')
                        break
                if buff.rstrip().endswith('$') or buff.find('denied')>=0:
                    s.close()
                    sys.exit(0)
                list_buff=buff.splitlines()
                list_buff.insert(-1,'\n已切换为root用户\n')
                buff=''.join(list_buff)
                print(buff,end='')

                #运行命令
                result += buff
                for cmd in cmds:
                    buff=''
                    ssh.send(cmd+'\n')
                    buff = ssh.recv(9999).decode('utf-8')
                    while not buff.rstrip().endswith('#'):
                        buff += ssh.recv(9999).decode('utf-8')
                    print(buff,end='')
                    result += buff

                s.close()
                # result = buff
                return result

        except paramiko.ssh_exception.AuthenticationException:
            print('Failed to login. ip username or password not correct.')
            try:
                client.close()
            except:
                pass
        except Exception as e:
            print("连接超时")
            try:
                client.close()
            except:
                print(e)
            # sys.exit(1)

    def upload_file(self,files=[],path='my_upload'):

        """向服务器上传多个文件
        Args:
            files:  本地文件的完整路径列表,默认为空列表.
            path:   上传文件的目的路径，默认在当前目录下创建my_upload
        """

        print('即将连接到%s' % self.host)
        print('@@@@@@@@@@@@')
        s = paramiko.SSHClient()
        s.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            s.connect(self.host, self.port, self.user, self.user_pwd,timeout=100)
            sftp=s.open_sftp()
        except paramiko.ssh_exception.AuthenticationException:
            print('Failed to login. ip username or password not correct.')
            exit(-1)

        # tran = paramiko.Transport(self.host,self.port)
        # tran.connect(self.user,self.user_pwd)
        # sftp = paramiko.SFTPClient.from_transport(tran)

        try:
            sftp.mkdir(path)
            print('上传目录{}创建成功'.format(path))
        except IOError:
            print('目录{}已存在'.format(path))
            
        sftp.chdir('./'+path)
        path_dir=sftp.getcwd()
        print("进入目录{}".format(path_dir))
        print("当前文件列表:\n{}".format(sftp.listdir('.')))
        if files:
            for file_local in files:
                file_name=os.path.basename(file_local)
                file_remote=path_dir+'/'+file_name
                result=sftp.put(file_local,file_remote,confirm=True)
                print('{}  {}'.format(result,file_name))
                # file_remote='/home/sftpuser/my_upload/1'
                # sftp.put(file_local,file_remote)
            print('*文件上传成功*')
            print("当前文件列表:\n{}".format(sftp.listdir('.')))

        s.close()

    def download_file(self,files=[],path='my_download'):
        print('即将连接到%s' % self.host)
        print('@@@@@@@@@@@@')
        s = paramiko.SSHClient()
        s.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            s.connect(self.host, self.port, self.user, self.user_pwd,timeout=100)
            sftp=s.open_sftp()
        except paramiko.ssh_exception.AuthenticationException:
            print('Failed to login. ip username or password not correct.')
            exit(-1)

        try:
            sftp.mkdir(path)
            print('远程下载目录{}创建成功'.format(path))
        except IOError:
            print('远程下载目录{}已存在'.format(path))
            sftp.rename(path,path+'_'+time.strftime('%Y%m%d%H%M%S'))
            sftp.mkdir(path)

        sftp.chdir('./'+path)
        path_dir=sftp.getcwd()
        s.close()

        print('断开连接%s, 开始初始化' % self.host)

        files_download=[]
        cmds=[]
        #创建下载目录的文件路径列表
        #将目的文件复制到下载目录my_download
        if files:
            for file in files:
                file_name=os.path.basename(file)
                file_download=path_dir+'/'+file_name
                files_download.append(file_download)
                cmd1=f'cp -f {file} {file_download}'
                cmd2=f'chmod 777 {file_download}'
                cmds.append(cmd1)
                cmds.append(cmd2)
            self.exec_command_root(cmds)

        print('初始化完成，重新连接到%s' % self.host)

        s = paramiko.SSHClient()
        s.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            s.connect(self.host, self.port, self.user, self.user_pwd,timeout=100)
            sftp=s.open_sftp()
        except paramiko.ssh_exception.AuthenticationException:
            print('Failed to login. ip username or password not correct.')
            exit(-1)

        if files_download:
            for file in files_download:
                file_name=os.path.basename(file)
                # file_local=path_dir+'/'+file_name
                file_local=os.path.join('C:\\Users\\10282787\\Downloads',file_name)
                # sftp.chmod(file,777)
                try:
                    sftp.get(file,file_local)
                    print('{}文件下载成功'.format(file_name))
                except FileNotFoundError:
                    print('没有文件',file)
                # file_remote='/home/sftpuser/my_upload/1'
                # sftp.put(file_local,file_remote)
            print("已下载文件列表:\n{}".format(os.listdir('C:\\Users\\10282787\\Downloads')))

        s.close()


if __name__ == '__main__':
    host="192.168.175.150"
    # host="192.168.0.1"
    port="22"
    user="tmpuser"
    # user_pwd="hello123"   # sftp user
    user_pwd="hello123"     # login user
    root_pwd="hello123"     # root user
    commands=['uptime','free -h','echo "hello world"']
    client = SshToClient(host, port, user, user_pwd, root_pwd)
    result = client.exec_command_root(commands)
    print('result:', result)
    #client.upload_file(files=['D:\software\VSCode\Project_study\paramiko_github\demo_simple.log','D:\itool_install\itooldevice-ott1-all.csv','D:\itool_install\itooldevice-ott2-all.csv'])
    # client.download_file(files=['/tmp/test/1.txt','/tmp/test/2.txt','/tmp/test/3.txt'])
